from model import Model
from openai import OpenAI
from zhipuai import ZhipuAI
from utils import LOG

class ModelFactory:
    @staticmethod
    def create_model(model_type, model_name, api_key):
        if model_type == 'OpenAIModel':
            return OpenAIModel(model_name, api_key)
        elif model_type == 'ChatGLMModel':
            return ChatGLMModel(model_name, api_key)
        else:
            raise ValueError(f'Invalid model type: {model_type}')


class OpenAIModel(Model):
    def __init__(self, model_name: str, api_key: str):
        super().__init__(model_name, api_key, OpenAI)


class ChatGLMModel(Model):
    def __init__(self, model_name: str, api_key: str):
        super().__init__(model_name, api_key, ZhipuAI)


if __name__ == '__main__':
    # Test Create OpenAIModel And Make Request
    model_type = 'OpenAIModel'
    model_name = 'gpt-3.5-turbo'
    api_key = 'your_openai_api_key'  # Replace your openai api key
    openai_model = ModelFactory.create_model(model_type, model_name, api_key)
    resp = openai_model.make_request('你好')
    print(resp)

    # Test Create ChatGLMModel And Make Request
    model_type = 'ChatGLMModel'
    model_name = 'glm-4'
    api_key = 'your_zhipuai_api_key'  # Replace your zhipuai api key
    chatglm_model = ModelFactory.create_model(model_type, model_name, api_key)
    resp = chatglm_model.make_request('你好')
    print(resp)